package com.autonavi.yunda.yunji.common.thread;

import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.slf4j.MDC;
import org.springframework.util.DigestUtils;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.*;

import static com.autonavi.yunda.yunji.common.constant.CommonConstant.MDC_TRACE_ID;

/**
 * @author cizhong.wcs
 * @date 2022/6/21 10:21
 */
@Slf4j
public class TraceableThreadPoolExecutor extends ThreadPoolExecutor {
    private final ThreadLocal<Long> startTime = new ThreadLocal<>();

    private int totalTasks;

    private int deltaTasks;

    private long deltaTime;

    public TraceableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue, @NotNull ThreadFactory threadFactory, @NotNull RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    @Override
    public void execute(@NotNull Runnable command) {
        super.execute(getTraceableRunnableWrapper(command));
    }


    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        super.beforeExecute(t, r);
        startTime.set(System.currentTimeMillis());
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        long time = System.currentTimeMillis() - startTime.get();
        synchronized (this) {
            deltaTime += time;
            deltaTasks++;
            totalTasks++;
        }
        super.afterExecute(r, t);
    }

    public synchronized int getTotalTasks() {
        return totalTasks;
    }

    public synchronized double getAverageTaskTime() {
        double time = (deltaTasks == 0) ? 0 : deltaTime / (double) deltaTasks;
        deltaTime = 0;
        deltaTasks = 0;
        return time;
    }

    private Runnable getTraceableRunnableWrapper(final Runnable command) {
        final Map<String, Object> contextMap = AmapThreadContext.getContextMapCopy();
        Map<String, String> mdcContextMap = MDC.getCopyOfContextMap();
        if (mdcContextMap == null) {
            mdcContextMap = new HashMap<>();
        }
        if (!mdcContextMap.containsKey(MDC_TRACE_ID)) {
            mdcContextMap.put(MDC_TRACE_ID, DigestUtils.md5DigestAsHex(UUID.randomUUID().toString().getBytes()));
        }
        final Map<String, String> finalMdcContextMap = mdcContextMap;
        return () -> {
            try {
                AmapThreadContext.initialThreadContext(contextMap, finalMdcContextMap);
                command.run();
                try {
                    if (command instanceof FutureTask) {
                        FutureTask<?> futureTask = (FutureTask<?>) command;
                        futureTask.get(500, TimeUnit.MILLISECONDS);
                    }
                } catch (ExecutionException e) {
                    log.warn("thread run error", e);
                } catch (InterruptedException | TimeoutException ee) {
                    log.warn("thread get task result failed.");
                }
            } finally {
                AmapThreadContext.clearContext();
            }
        };
    }
}
